import cv2
import matplotlib.pyplot as plt
from utils import show_det_result
import numpy as np
import mxnet as mx
from model import SSD
import time
from mxnet import ndarray as nd
from mxnet.ndarray.contrib import MultiBoxDetection
from mxnet import image
std = np.array([58.395, 57.12, 57.375])
rgb_mean = np.array([130.063048, 129.967301, 124.410760])
ctx = mx.gpu(0)
resize = (512, 512)
num_class = 1
class_names = ['papercup']
# 加载模型

def loadWeight():
    net = SSD(num_class, ctx=ctx, verbose=False, prefix="ssd_")
    net.load_params("models/sdl_coin_vgg11bn29_512x512_data_sizes.param", ctx=ctx)
    return net
# 读入图片

def predict(img_nd, net):
    # predict
    tic = time.time()
    anchors, box_preds, cls_preds = net(img_nd)
    # 处理结果
    cls_probs = nd.SoftmaxActivation(cls_preds.transpose((0, 2, 1)), mode='channel')
    out = MultiBoxDetection(cls_probs, box_preds, anchors, force_suppress=True, clip=False, nms_threshold=0.2)

    out = out.asnumpy()
    print(out.shape)
    print('FORWARD TIME', time.time() - tic)
    return out
def detector(net, img_paths, threshold = 0.33):
    img_nds = None
    print(img_paths)
    tic = time.time()
    for img_path in img_paths:
        # 读入图片
        # img = image.imread(img_path)

        img = plt.imread(img_path)
        img = (cv2.resize(img, resize) - rgb_mean) / std
        img_nd = nd.array(img, ctx=ctx)
        img_nd = img_nd.expand_dims(0).transpose((0, 3, 1, 2))
        if img_nds is None:
            img_nds = img_nd
        else:
            img_nds = nd.concat(img_nds, img_nd, dim = 0)
    print('IO TIME:', time.time() - tic)
    outs = predict(img_nds, net)

    all_results = []
    for out in outs:
        results = []
        colom_mask = (out[: ,1] > threshold) * (out[:, 0] != -1)
        out = out[colom_mask, :]
        for item in out:
            class_name = class_names[int(item[0])]
            prob = float(item[1])
            cx = float((item[2] + item[4]) / 2)
            cy = float((item[3] + item[5]) / 2)
            w = float((item[4] - item[2]))
            h = float((item[5] - item[3]))
            result = [class_name, prob, [cx, cy, w, h]]
            results.append(result)
        all_results.append(results)
    return all_results



# 绘制图片

# imgs = (img_nds.transpose((0, 2, 3, 1)).asnumpy() * std)+ rgb_mean
# show_det_result(imgs[0], out[0], threshold=0.4)
if __name__ == '__main__':
    img_paths = ['test.jpg',] * 8
    net = loadWeight()
    outs = detector(net, img_paths)
    print(outs)
    # for i, out in enumerate(outs):
    #     _, figs = plt.subplots()
    #     img = plt.imread(img_paths[i])
    #     figs.imshow(img)
    #     tmp = [img.shape[1], img.shape[0]] * 2
    #     for item in out:
    #         box = np.array(item[2]) * tmp
    #         rect = plt.Rectangle((box[0] - box[2] / 2, box[1] - box[3] / 2), box[2], box[3], fill=False, color='blue')
    #         figs.add_patch(rect)
    #     plt.show()
    # print(outs)